from llmtcg.chat.model import ChatModel, ChatContext
from llmtcg.chat.role import Role

from loguru import logger

from anthropic import Anthropic
from anthropic import PermissionDeniedError, RateLimitError, APIConnectionError, InternalServerError
from anthropic.types.message import Message
from anthropic.types.message_param import MessageParam

from dataclasses import dataclass, field
from typing import Optional, List, Tuple
from os import environ
from enum import Enum
from time import sleep

from llmtcg.utils.profiler import Currency


@dataclass
class ClaudeChatModel(ChatModel):
    """ The set of anthropic models deployed at its official website. """

    class Model(Enum):
        OPUS = "claude-3-opus-20240229" # claude-opus-4-20250514
        SONNET = "claude-3-sonnet-20240229" # claude-sonnet-4-20250514
        HAIKU = "claude-3-haiku-20240307" # claude-3-5-haiku-20241022

    api_key: str = field(default=environ.get("ANTHROPIC_API_KEY", ""))
    model: 'ClaudeChatModel.Model' = Model.SONNET

    """ The connection client to the claudi API """
    _client: Optional[Anthropic] = field(default=None, init=False)

    def pretty_name(self) -> str:
        return "claude-" + self.model.name.lower()

    @property
    def client(self) -> Anthropic:
        if self._client is not None:
            return self._client

        conn = Anthropic(api_key=self.api_key)
        self._client = conn
        return conn

    @staticmethod
    def get_plain_text(message: Message) -> str:
        text = ""
        for blk in message.content:
            text += blk.text

        return text

    def chat(self, context: ChatContext) -> str:
        history_messages: List[MessageParam] = []
        for role, msg in context.history:
            history_messages.append(MessageParam(
                role=role.value,  # type: ignore
                content=msg
            ))

        def send_request() -> Message:
            message: Message = self.client.messages.create(
                model=self.model.value,
                max_tokens=4096,
                temperature=0,
                system=context.system_prompt,
                messages=history_messages
            )
            return message

        while True:
            try:
                message = send_request()
                break
            except RateLimitError:
                logger.warning("rate limitation reached, will retry after 10 seconds")
                sleep(10)
            except APIConnectionError:
                logger.warning("api connection failed, will retry after 5 seconds")
                sleep(5)
            except PermissionDeniedError:
                logger.warning("encountered permission denied issue, will retry after 10 seconds")
                sleep(10)
            except InternalServerError:
                logger.warning("encountered internal server error, will retry after 10 seconds")
                sleep(10)

        response_text = self.get_plain_text(message)
        self.notify_usage(message.usage.input_tokens, message.usage.output_tokens)
        return response_text

    def price(self) -> Tuple[float, float, Currency]:
        if self.model is self.Model.HAIKU:
            return 0.25, 1.25, Currency.USD
        elif self.model is self.Model.SONNET:
            return 3.00, 15.00, Currency.USD
        elif self.model is self.Model.OPUS:
            return 15.00, 75.00, Currency.USD

        raise NotImplementedError(self.model)

    def __hash__(self) -> int:
        return id(self)
